"""
python fujian1_arimaAndOrigin_gragh.py
"""

import os
import json
import matplotlib.pyplot as plt
import pandas as pd

def plot_inventory_data(original_data_dir, processed_data_dir, output_dir):
    """
    画图，展示每个类别的原始数据和处理后的数据。

    参数:
    - original_data_dir: str, 原始数据的 JSON 文件夹路径
    - processed_data_dir: str, 处理后数据的 JSON 文件夹路径
    - output_dir: str, 输出图形的文件夹路径
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有原始数据文件
    for original_file in os.listdir(original_data_dir):
        if original_file.endswith('.json'):
            category_name = original_file.split('_')[1].split('.')[0]  # 从文件名中提取类别
            
            # 读取原始数据
            with open(os.path.join(original_data_dir, original_file), 'r') as f:
                original_data = json.load(f)
            
            # 读取处理后的数据
            processed_file = f'smoothAfterSpline_category{category_name}.json'
            with open(os.path.join(processed_data_dir, processed_file), 'r') as f:
                processed_data = json.load(f)

            # 转换为 DataFrame
            original_df = pd.DataFrame(original_data)
            processed_df = pd.DataFrame(processed_data)

            # 将日期列转换为 datetime 类型
            original_df['date'] = pd.to_datetime(original_df['date'])
            processed_df['date'] = pd.to_datetime(processed_df['date'])

            # 绘制图形
            plt.figure(figsize=(12, 6))
            plt.plot(original_df['date'], original_df['inventory'], 'o-', label='原始数据 (12个月)', color='blue')
            plt.plot(processed_df['date'], processed_df['inventory'], 's-', label='处理后数据 (15个月)', color='orange')

            # 添加图例和标题
            plt.title(f'库存数据对比 - 类别 {category_name}')
            plt.xlabel('日期')
            plt.ylabel('库存')
            plt.xticks(rotation=45)
            plt.legend()
            plt.tight_layout()

            # 保存图形
            plt.savefig(os.path.join(output_dir, f'category_{category_name}_inventory_plot.png'))
            plt.close()  # 关闭当前图形以释放内存

if __name__ == "__main__":
    original_data_dir = r"..\fujian\fujian1\cubic_spline\interpolation_out"
    processed_data_dir = r"..\fujian\fujian1\pmdarima_single_batch"
    output_dir = r"..\fujian\fujian1\pmdarima_single_batch\plots"
    
    plot_inventory_data(original_data_dir, processed_data_dir, output_dir)
